﻿#include "CCaptureProcess.h"
int find_video_stream_index(const AVFormatContext* pFormatContext)
{
    int nIndex = -1;

    for (int i = 0; i < pFormatContext->nb_streams; ++i)
    {
        if (AVMEDIA_TYPE_VIDEO == pFormatContext->streams[i]->codec->codec_type)
        {
            nIndex = i;
            break;
        }
    }

    return nIndex;
}

uint8_t* getFrameBuffer(AVPicture* pFrame, const int destWidth, const int destHeight)
{
    int width = destWidth;
    int height = destHeight;

    uint8_t* buffer = (uint8_t*)malloc(height * width * 3);
    for (int y = 0; y < height; y++)
    {
        memcpy(buffer + y * pFrame->linesize[0], pFrame->data[0] + y * pFrame->linesize[0], width * 3);
    }

    return buffer;
}

void execConvert(AVFormatContext* pFormatContext, AVCodecContext* pCodecContext,
    int nStreamsIndex, vector<ImageData*>& vecImageDatas, int destWidth, int destHeight)
{
    float scale = (float)destWidth / pCodecContext->width;
    destHeight = pCodecContext->height * scale;
    //分配视频帧
    //确定需要的缓冲区大小，并分配缓冲区
    int numBytes = avpicture_get_size(AV_PIX_FMT_RGB24, destWidth,
        destHeight);

    uint8_t* buffer = (uint8_t*)av_malloc(numBytes * sizeof(uint8_t));

    AVPicture* pFrameRGB = (AVPicture*)av_frame_alloc();
    avpicture_fill(pFrameRGB, buffer, AV_PIX_FMT_RGB24, destWidth, destHeight);


    //初始化sws上下文
    SwsContext* sws_ctx = sws_getContext(pCodecContext->width, pCodecContext->height,
        pCodecContext->pix_fmt, destWidth, destHeight, AV_PIX_FMT_RGB24, SWS_BILINEAR,
        nullptr, nullptr, nullptr);

    //读取数据包中的视频流，并将其解码到帧中
    AVPacket packet;
    AVFrame* pFrame = av_frame_alloc();

    while (av_read_frame(pFormatContext, &packet) >= 0)
    {
        //判断是不是来自视频流的数据包
        if (packet.stream_index == nStreamsIndex)
        {
            int frameFinished = 0;
            //解码视频流
            avcodec_decode_video2(pCodecContext, pFrame, &frameFinished, &packet);
            //是否得到一个视频帧
            if (frameFinished)
            {
                //将图像转换为RGB
                sws_scale(sws_ctx, pFrame->data,
                    pFrame->linesize, 0, pCodecContext->height, pFrameRGB->data, pFrameRGB->linesize);
                
                //存储帧图
                ImageData* imageData = (ImageData *)malloc(sizeof(ImageData));
                imageData->width = (uint32_t)destWidth;
                imageData->height = (uint32_t)destHeight;

                imageData->data = getFrameBuffer(pFrameRGB, destWidth, destHeight);

                vecImageDatas.push_back(imageData);
            }
        }
    }

    //释放内存
    if (nullptr != pFrame)
    {
        av_free(pFrame);
        pFrame = nullptr;
    }

    if (nullptr != sws_ctx)
    {
        sws_freeContext(sws_ctx);
        sws_ctx = nullptr;
    }

    if (nullptr != pFrameRGB)
    {
        av_free(pFrameRGB);
        pFrameRGB = nullptr;
    }

    if (nullptr != buffer)
    {
        av_free(buffer);
        buffer = nullptr;
    }
}

void captureProcess(AVFormatContext* pFormatCtx, vector<ImageData*>& vecImageDatas, int width, int height)
{
    if (nullptr == pFormatCtx)
    {
        return;
    }

    //检查视频文件中的流信息
    if (avformat_find_stream_info(pFormatCtx, nullptr) < 0)
    {
        return;
    }

    //遍历数组找到视频流索引
    int nStreamsIndex = find_video_stream_index(pFormatCtx);

    //获取指向视频流的上下文指针
    AVCodecContext* pCodecContext = pFormatCtx->streams[nStreamsIndex]->codec;

    //找到视频流的编解码器
    AVCodec* pCodec = avcodec_find_decoder(pCodecContext->codec_id);
    if (nullptr == pCodec)
    {
        avcodec_close(pCodecContext);
        return;
    }

    //打开编解码器
    if (avcodec_open2(pCodecContext, pCodec, nullptr) < 0)
    {
        avcodec_close(pCodecContext);
        return;
    }

    //执行转换
    execConvert(pFormatCtx, pCodecContext, nStreamsIndex, vecImageDatas, width, height);

    avcodec_close(pCodecContext);
}

